import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Polygon, FancyBboxPatch

# ---------------- 全局配置 ----------------
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Times New Roman"]
plt.rcParams["axes.unicode_minus"] = False
plt.rcParams.update({
    "font.size": 12,
    "axes.labelsize": 14,
    "axes.titlesize": 16
})

# 尝试开启 LaTeX
try:
    plt.rcParams["text.usetex"] = True
    plt.rcParams["text.latex.preamble"] = [
        r"\usepackage{amsmath}",
        r"\usepackage{amsfonts}",
        r"\usepackage{amsbsy}"
    ]
except Exception:
    plt.rcParams["text.usetex"] = False


def draw_arrow(ax, x1, y1, x2, y2, width=1.2, color='black', style='->'):
    """绘制箭头"""
    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle=style, color=color, lw=width))


def plot_improved_pso_flowchart():
    """绘制改进PSO求解流程图"""

    # ---------------- 基本设置 ----------------
    fig, ax = plt.subplots(figsize=(14, 10))
    ax.set_xlim(0, 12)
    ax.set_ylim(0, 10)
    ax.axis('off')

    box_width = 6
    box_height = 1.2
    step_color = '#e6f7ff'       # 浅蓝
    decision_color = '#ffe6cc'   # 浅橙
    text_color = 'black'

    # ---------------- 流程框 ----------------
    # 1. 粒子编码与初始化
    x1, y1 = 3, 8.5
    ax.add_patch(FancyBboxPatch((x1, y1), box_width, box_height,
                                boxstyle="round,pad=0.2", facecolor=step_color,
                                edgecolor="black", linewidth=2))
    ax.text(x1 + box_width/2, y1 + box_height/2,
            r"\textbf{① 粒子编码与初始化}" + "\n" +
            r"输入: 决策变量范围" + "\n" +
            r"输出: 初始粒子群 $\{x_i, v_i\}$",
            ha='center', va='center', fontsize=12, color=text_color)

    # 2. 适应度函数计算
    x2, y2 = 3, 6.5
    ax.add_patch(FancyBboxPatch((x2, y2), box_width, box_height,
                                boxstyle="round,pad=0.2", facecolor=step_color,
                                edgecolor="black", linewidth=2))
    ax.text(x2 + box_width/2, y2 + box_height/2,
            r"\textbf{② 适应度函数计算}" + "\n" +
            r"(含约束验证、遮蔽判定)" + "\n" +
            r"输入: 粒子群 $\{x_i\}$" + "\n" +
            r"输出: 适应度值 $\{f(x_i)\}$",
            ha='center', va='center', fontsize=12, color=text_color)

    # 3. 粒子速度与位置更新
    x3, y3 = 3, 4.5
    ax.add_patch(FancyBboxPatch((x3, y3), box_width, box_height,
                                boxstyle="round,pad=0.2", facecolor=step_color,
                                edgecolor="black", linewidth=2))
    ax.text(x3 + box_width/2, y3 + box_height/2,
            r"\textbf{③ 粒子速度与位置更新}" + "\n" +
            r"(含自适应惯性权重)" + "\n" +
            r"输入: $\{x_i, v_i, p_i, g\}$" + "\n" +
            r"输出: 更新后 $\{x_i', v_i'\}$",
            ha='center', va='center', fontsize=12, color=text_color)

    # 4. 精英学习
    x4, y4 = 3, 2.5
    ax.add_patch(FancyBboxPatch((x4, y4), box_width, box_height,
                                boxstyle="round,pad=0.2", facecolor=step_color,
                                edgecolor="black", linewidth=2))
    ax.text(x4 + box_width/2, y4 + box_height/2,
            r"\textbf{④ 精英学习 (局部扰动)}" + "\n" +
            r"输入: 全局最优解 $g$" + "\n" +
            r"输出: 优化后 $g'$",
            ha='center', va='center', fontsize=12, color=text_color)

    # 5. 收敛判定（菱形）
    x5, y5 = 3, 0.5
    diamond = Polygon([[x5 + box_width/2, y5 + box_height],
                       [x5 + box_width, y5 + box_height/2],
                       [x5 + box_width/2, y5],
                       [x5, y5 + box_height/2]],
                      facecolor=decision_color, edgecolor='black', linewidth=2)
    ax.add_patch(diamond)
    ax.text(x5 + box_width/2, y5 + box_height/2,
            r"\textbf{⑤ 收敛判定?}" + "\n" +
            r"(迭代次数 / 精度)" + "\n" +
            r"输入: 当前迭代状态" + "\n" +
            r"输出: 判定结果",
            ha='center', va='center', fontsize=12, color=text_color)

    # 6. 输出最优解
    x6, y6 = 10.5, 0.5
    ax.add_patch(FancyBboxPatch((x6 - box_width/2, y6), box_width, box_height,
                                boxstyle="round,pad=0.2", facecolor=step_color,
                                edgecolor="black", linewidth=2.5))
    ax.text(x6, y6 + box_height/2,
            r"\textbf{输出最优解}" + "\n" +
            r"$x^* = \arg\min f(x)$" + "\n" +
            r"$f(x^*)$",
            ha='center', va='center', fontsize=12, color=text_color)

    # ---------------- 箭头 ----------------
    draw_arrow(ax, x1 + box_width/2, y1, x2 + box_width/2, y2 + box_height)
    draw_arrow(ax, x2 + box_width/2, y2, x3 + box_width/2, y3 + box_height)
    draw_arrow(ax, x3 + box_width/2, y3, x4 + box_width/2, y4 + box_height)
    draw_arrow(ax, x4 + box_width/2, y4, x5 + box_width/2, y5 + box_height)

    # 是分支 5->6
    draw_arrow(ax, x5 + box_width, y5 + box_height/2, x6 - box_width/2, y6 + box_height/2)
    ax.text(x5 + box_width + 0.3, y5 + box_height/2, "是",
            ha='left', va='center', fontsize=12, color="red")

    # 否分支 5->2
    mid_x = x5 - 1
    mid_y1 = y5 - 0.5
    mid_y2 = y2 + box_height + 0.5
    draw_arrow(ax, x5, y5 + box_height/2, mid_x, mid_y1)
    draw_arrow(ax, mid_x, mid_y1, mid_x, mid_y2)
    draw_arrow(ax, mid_x, mid_y2, x2 + box_width/2, mid_y2)
    draw_arrow(ax, x2 + box_width/2, mid_y2, x2 + box_width/2, y2 + box_height)
    ax.text(mid_x - 0.3, (mid_y1 + mid_y2)/2, "否",
            ha='right', va='center', fontsize=12, color="blue")

    # ---------------- 说明与标题 ----------------
    plt.title("改进PSO求解流程图", fontsize=20, pad=20, weight="bold")

    plt.figtext(0.5, 0.01,
                ("图注：以流程图形式展示改进PSO求解过程：\n"
                 "① 粒子编码与初始化 → ② 适应度函数计算（含约束验证、遮蔽判定） →\n"
                 "③ 粒子速度与位置更新（含自适应惯性权重） → ④ 精英学习（局部扰动） →\n"
                 "⑤ 收敛判定（是则输出最优解，否则返回步骤②），"
                 "并标注各步骤的核心输入与输出。"),
                ha="center", fontsize=12,
                bbox=dict(facecolor="white", alpha=0.8, boxstyle="round,pad=0.5"))

    plt.tight_layout()
    plt.savefig("improved_pso_flowchart.png", dpi=300, bbox_inches="tight")
    plt.show()


if __name__ == "__main__":
    plot_improved_pso_flowchart()
